from langchain_openai import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
import os
import re
from dotenv import load_dotenv

# 加载环境变量（从项目根目录的.env文件）
load_dotenv()

# 配置参数
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
# 由于网络的原因这里我使用的是openai代理地址和key
OPENAI_API_BASE = os.getenv("OPENAI_API_BASE", "https://api.openai-proxy.org/v1")
MD_FILE_PATH = "/Users/wu/PycharmProjects/PythonProject/output/spring-ai/1.0/content/结构化输出.md"

# 检查必要的环境变量
if not OPENAI_API_KEY:
    raise ValueError("请在项目根目录的.env文件中设置 OPENAI_API_KEY")

def count_words_and_time(text):
    """
    统计文本字数和预估阅读时间
    Args:
        text: 文本内容
    Returns:
        tuple: (字数, 阅读时间分钟数)
    """
    # 移除markdown语法和空白字符
    clean_text = re.sub(r'```.*?```', '', text, flags=re.DOTALL)  # 移除代码块
    clean_text = re.sub(r'`.*?`', '', clean_text)  # 移除行内代码
    clean_text = re.sub(r'\[.*?\]\(.*?\)', '', clean_text)  # 移除链接
    clean_text = re.sub(r'[#*_~]', '', clean_text)  # 移除markdown符号
    clean_text = re.sub(r'\s+', '', clean_text)  # 移除空白字符
    
    # 统计中文字符和英文单词
    chinese_chars = len(re.findall(r'[\u4e00-\u9fff]', clean_text))
    english_words = len(re.findall(r'[a-zA-Z]+', clean_text))
    
    # 计算总字数（中文字符 + 英文单词）
    total_words = chinese_chars + english_words
    
    # 预估阅读时间（假设中文阅读速度300字/分钟，英文阅读速度200词/分钟）
    reading_time = (chinese_chars / 300) + (english_words / 200)
    
    return total_words, round(reading_time, 1)

def detect_language(text):
    """
    通过分析markdown文件中的标题来判断文档语言
    Args:
        text: markdown文件内容
    Returns:
        str: 语言代码 ('zh' 或 'en')
    """
    # 提取所有markdown标题（从#到######）
    headers = re.findall(r'^#{1,6}\s+(.+)$', text, re.MULTILINE)
    
    if not headers:
        return 'en'  # 如果没有找到标题，默认使用英文
    
    # 统计标题中的中文字符
    chinese_chars = 0
    total_chars = 0
    
    for header in headers:
        # 清理标题中的markdown格式（如*、_、~等）
        header = re.sub(r'[`*_~]', '', header)
        chinese_chars += len(re.findall(r'[\u4e00-\u9fff]', header))
        total_chars += len(re.sub(r'\s', '', header))
    
    print(f"找到 {len(headers)} 个标题")
    print(f"标题中的中文字符数: {chinese_chars}")
    print(f"标题中的总字符数: {total_chars}")
    
    # 如果标题中超过30%是中文，则判定为中文文档
    if total_chars > 0 and chinese_chars / total_chars > 0.3:
        return 'zh'
    return 'en'

def get_language_prompt(lang):
    """
    根据语言获取对应的提示语
    Args:
        lang: 语言代码
    Returns:
        str: 对应语言的提示语
    """
    prompts = {
        'en': "You are a helpful assistant that creates concise summaries of markdown content in English. "
              "Focus on the main points and key information while maintaining the essential context.",
        'zh': "你是一个帮助生成中文markdown内容摘要的助手。"
              "请关注主要内容要点和关键信息，同时保持必要的上下文。",
        'ja': "あなたはmarkdownコンテンツの簡潔な要約を作成するアシスタントです。"
              "主要なポイントと重要な情報に焦点を当てながら、必要な文脈を維持してください。",
        'ko': "당신은 markdown 콘텐츠의 간결한 요약을 생성하는 도우미입니다. "
              "주요 포인트와 핵심 정보에 초점을 맞추면서 필수적인 맥락을 유지하세요."
    }
    return prompts.get(lang, prompts['en'])  # 如果语言不支持，默认使用英文

def insert_summary_after_front_matter(content, summary, word_count, reading_time, lang):
    """
    在YAML front matter后插入摘要
    Args:
        content: 原始markdown内容
        summary: 生成的摘要
        word_count: 字数统计
        reading_time: 阅读时间
        lang: 语言代码
    Returns:
        str: 插入摘要后的完整内容
    """
    # 根据语言选择统计信息的显示文本
    stats_text = {
        'zh': f"\n\n### 文档统计\n- 总字数：{word_count} 字\n- 预计阅读时间：{reading_time} 分钟\n",
        'en': f"\n\n### Document Statistics\n- Total Words: {word_count}\n- Estimated Reading Time: {reading_time} minutes\n"
    }.get(lang, f"\n\n### Document Statistics\n- Total Words: {word_count}\n- Estimated Reading Time: {reading_time} minutes\n")

    # 查找front matter的结束位置（第二个---标记）
    front_matter_end = content.find('---', content.find('---') + 3)
    if front_matter_end == -1:
        # 如果没有找到front matter，则在文件开头插入摘要
        return f"## Summary\n{summary}{stats_text}\n\n{content}"
    
    # 在front matter后插入摘要
    return content[:front_matter_end + 3] + f"\n\n## Summary\n{summary}{stats_text}\n\n" + content[front_matter_end + 3:]

def summarize_markdown(md_file_path):
    """
    生成markdown文件的摘要并插入到文件中
    Args:
        md_file_path: markdown文件路径
    Returns:
        bool: 操作是否成功
    """
    # 初始化OpenAI模型
    llm = ChatOpenAI(
        model="gpt-3.5-turbo",
        temperature=0.7,
        api_key=OPENAI_API_KEY,
        base_url=OPENAI_API_BASE
    )

    try:
        # 读取markdown文件
        with open(md_file_path, 'r', encoding='utf-8') as file:
            content = file.read()

        # 检测文档语言
        lang = detect_language(content)
        print(f"检测到的语言: {lang}")

        # 统计字数和阅读时间
        word_count, reading_time = count_words_and_time(content)
        print(f"文档字数: {word_count}")
        print(f"预计阅读时间: {reading_time} 分钟")

        # 创建对应语言的提示模板
        prompt = ChatPromptTemplate.from_messages([
            ("system", get_language_prompt(lang)),
            ("user", "Please summarize the following markdown content:\n\n{text}")
        ])

        # 创建处理链
        chain = prompt | llm

        # 生成摘要
        summary = chain.invoke({"text": content})

        # 在front matter后插入摘要和统计信息
        new_content = insert_summary_after_front_matter(content, summary.content, word_count, reading_time, lang)
        with open(md_file_path, 'w', encoding='utf-8') as file:
            file.write(new_content)

        print(f"摘要和统计信息已成功添加到 {md_file_path}")
        return True

    except Exception as e:
        print(f"处理文件时出错: {str(e)}")
        return False

if __name__ == "__main__":
    if os.path.exists(MD_FILE_PATH):
        summarize_markdown(MD_FILE_PATH)
    else:
        print("文件未找到！")
